#include "Encoder.h"

 Encoder::Encoder(QObject *parent):QObject(parent)
 {

 }
 Encoder::~Encoder()
 {
    av_frame_free(&frame);
    avcodec_free_context(&codecCtx);
    sws_freeContext(swsCtx);
 }
int  Encoder::initialize(int width, int height,int fps)
{   
    
    m_width = width;
    m_height = height;
    m_fps = fps;
    
    // 查找H.264编码器
    codec = (AVCodec *)avcodec_find_encoder(AV_CODEC_ID_H264);
    if (!codec) {
        fprintf(stderr, "Codec not found\n");
        return -1;
    }
    // 分配编码器上下文
    codecCtx = avcodec_alloc_context3(codec);
    if (!codecCtx) {
        fprintf(stderr, "Could not allocate video codec context\n");
        return -1;
    }
    // 设置编码参数
    codecCtx->bit_rate = 400000;                 // 比特率
    codecCtx->width = width;                     // 宽度
    codecCtx->height = height;                   // 高度
    codecCtx->time_base = (AVRational){1, fps};  // 帧率
    codecCtx->framerate = (AVRational){fps, 1};  // 帧率
    codecCtx->gop_size = 10;                     // 每10帧插入一个I帧
    codecCtx->max_b_frames = 1;                  // 使用B帧
    codecCtx->pix_fmt = AV_PIX_FMT_YUV420P;      // 输出格式是YUV420P

    // 打开编码器
    if (avcodec_open2(codecCtx, codec, NULL) < 0) {
        fprintf(stderr, "Could not open codec\n");
        return -1;
    }

    // 创建帧并设置参数
    frame = av_frame_alloc();
    if (!frame) {
        fprintf(stderr, "Could not allocate video frame\n");
        return -1;
    }
    frame->format = codecCtx->pix_fmt;
    frame->width = codecCtx->width;
    frame->height = codecCtx->height;
    
    if (av_frame_get_buffer(frame, 32) < 0) {
        fprintf(stderr, "Could not allocate frame buffer\n");
        return -1;
    }

    // 创建SwsContext来转换像素格式（从RGBA到YUV420P）
    swsCtx = sws_getContext(width, height, AV_PIX_FMT_BGRA,
                                        width, height, AV_PIX_FMT_YUV420P,
                                        SWS_BILINEAR, NULL, NULL, NULL);


    // 初始化数据包
    av_init_packet(&pkt);
    pkt.data = NULL;
    pkt.size = 0;

    return 0;
}

int  Encoder::encodeFrame(QImage& img,int pts,VideoFrame& videoframe)
{
    // 将RGBA转换为YUV420P
    uint8_t *srcSlice[] = {img.bits(), nullptr};
    int srcStride[] = {img.width() * 4};
    sws_scale(swsCtx, srcSlice, srcStride, 0, img.height(), frame->data, frame->linesize);

    // 设置PTS（显示时间戳）
    frame->pts = pts;

    // 发送帧给编码器
    if (avcodec_send_frame(codecCtx, frame) < 0)
    {
        fprintf(stderr, "Error sending frame for encoding\n");
        return -1;
    }

    // 接收编码后的数据包
    while (avcodec_receive_packet(codecCtx, &pkt) == 0)
    {
        videoframe.data = QByteArray(reinterpret_cast<char *>(pkt.data), pkt.size);
        videoframe.size = QSize(codecCtx->width, codecCtx->height);
        av_packet_unref(&pkt);
    }

    return 0;
}

void Encoder::finish()
{
    // 刷新编码器（将剩余的数据编码并输出）
    avcodec_send_frame(codecCtx, NULL);
    while (avcodec_receive_packet(codecCtx, &pkt) == 0) {
        // fwrite(pkt.data, 1, pkt.size, file);
        av_packet_unref(&pkt);
    }
}